#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
训练监控组件 - 用于显示训练进度和结果
"""

import os
import numpy as np
from PyQt5.QtWidgets import (
    QWidget, QVBoxLayout, QHBoxLayout, QLabel, QProgressBar, 
    QPushButton, QTabWidget, QTableWidget, QTableWidgetItem,
    QComboBox, QGridLayout, QGroupBox, QFormLayout, QSpinBox,
    QDoubleSpinBox, QCheckBox, QHeaderView, QSplitter
)
from PyQt5.QtCore import Qt, pyqtSignal, QTimer
from PyQt5.QtGui import QColor
import threading

# 尝试导入matplotlib
try:
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
    MATPLOTLIB_AVAILABLE = True
except ImportError:
    MATPLOTLIB_AVAILABLE = False

class TrainingChart(QWidget):
    """训练图表组件"""
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__(parent)
        
        # 添加线程锁
        self.chart_lock = threading.Lock()
        
        # 检查matplotlib是否可用
        if not MATPLOTLIB_AVAILABLE:
            # 显示错误信息
            layout = QVBoxLayout(self)
            label = QLabel("缺少matplotlib库，无法显示图表")
            label.setAlignment(Qt.AlignCenter)
            layout.addWidget(label)
            return
        
        # 设置布局
        self.layout = QVBoxLayout(self)
        
        # 创建图表选择下拉框
        self.chart_selector = QComboBox()
        self.chart_selector.addItems(["损失", "准确率", "学习率"])
        self.chart_selector.currentIndexChanged.connect(self.update_chart)
        self.layout.addWidget(self.chart_selector)
        
        # 创建matplotlib图表
        self.figure = Figure(figsize=(5, 4), dpi=100)
        self.canvas = FigureCanvas(self.figure)
        self.layout.addWidget(self.canvas)
        
        # 初始化数据
        self.train_loss = []
        self.val_loss = []
        self.train_acc = []
        self.val_acc = []
        self.learning_rates = []
        self.epochs = []
        
        # 创建子图
        self.ax = self.figure.add_subplot(111)
    
    def update_data(self, epoch, metrics):
        """
        更新训练数据
        
        Args:
            epoch (int): 当前轮次
            metrics (dict): 训练指标字典，包含train_loss, val_loss, train_acc, val_acc等
        """
        # 使用互斥锁保护数据更新
        if hasattr(self, 'chart_lock') and self.chart_lock.acquire(False):
            try:
                # 添加轮次
                if epoch not in self.epochs:
                    self.epochs.append(epoch)
                
                # 更新损失
                if 'train_loss' in metrics:
                    while len(self.train_loss) <= epoch:
                        self.train_loss.append(None)
                    self.train_loss[epoch] = metrics['train_loss']
                
                if 'val_loss' in metrics:
                    while len(self.val_loss) <= epoch:
                        self.val_loss.append(None)
                    self.val_loss[epoch] = metrics['val_loss']
                
                # 更新准确率
                if 'train_acc' in metrics:
                    while len(self.train_acc) <= epoch:
                        self.train_acc.append(None)
                    self.train_acc[epoch] = metrics['train_acc']
                
                if 'val_acc' in metrics:
                    while len(self.val_acc) <= epoch:
                        self.val_acc.append(None)
                    self.val_acc[epoch] = metrics['val_acc']
                
                # 更新学习率
                if 'learning_rate' in metrics:
                    while len(self.learning_rates) <= epoch:
                        self.learning_rates.append(None)
                    self.learning_rates[epoch] = metrics['learning_rate']
                
                # 重绘图表
                # 在锁内调用更新图表，但使用QTimer延迟执行
                QTimer.singleShot(0, self.update_chart)
            finally:
                self.chart_lock.release()
        else:
            # 如果无法获取锁，跳过此次更新
            pass
    
    def update_chart(self):
        """更新当前图表"""
        if not MATPLOTLIB_AVAILABLE or not self.epochs:
            return
        
        # 使用互斥锁保护绘图操作
        if hasattr(self, 'chart_lock') and self.chart_lock.acquire(False):
            try:
                # 清除当前图表
                self.ax.clear()
                
                # 获取当前选择的图表类型
                chart_type = self.chart_selector.currentText()
                
                # 绘制选择的图表
                if chart_type == "损失":
                    self.ax.plot(self.epochs, self.train_loss, 'b-', label='训练损失')
                    self.ax.plot(self.epochs, self.val_loss, 'r-', label='验证损失')
                    self.ax.set_ylabel('损失')
                elif chart_type == "准确率":
                    self.ax.plot(self.epochs, self.train_acc, 'b-', label='训练准确率')
                    self.ax.plot(self.epochs, self.val_acc, 'r-', label='验证准确率')
                    self.ax.set_ylabel('准确率 (%)')
                elif chart_type == "学习率":
                    self.ax.plot(self.epochs, self.learning_rates, 'g-', label='学习率')
                    self.ax.set_ylabel('学习率')
                
                # 设置图表标题和轴标签
                self.ax.set_title(chart_type)
                self.ax.set_xlabel('轮次')
                
                # 添加图例
                self.ax.legend()
                
                # 添加网格
                self.ax.grid(True)
                
                # 重绘画布
                self.canvas.draw()
            finally:
                self.chart_lock.release()
        else:
            # 如果无法获取锁，跳过此次更新
            pass


class ConfusionMatrix(QWidget):
    """混淆矩阵组件"""
    
    # 添加信号
    cell_clicked = pyqtSignal(int, int)  # 发送单元格坐标 (真实类, 预测类)
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__(parent)
        
        # 设置布局
        self.layout = QVBoxLayout(self)
        
        # 创建标签
        self.title_label = QLabel("混淆矩阵")
        self.title_label.setAlignment(Qt.AlignCenter)
        self.layout.addWidget(self.title_label)
        
        # 创建表格
        self.table = QTableWidget(0, 0)
        self.table.setEditTriggers(QTableWidget.NoEditTriggers)
        self.table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.table.verticalHeader().setSectionResizeMode(QHeaderView.Stretch)
        self.layout.addWidget(self.table)
        
        # 存储标签列表
        self.labels = []
        self.matrix = None
        
        # 添加提示信息
        self.hint_label = QLabel("点击混淆矩阵单元格查看对应的问题样本")
        self.hint_label.setAlignment(Qt.AlignCenter)
        self.layout.addWidget(self.hint_label)
        
        # 连接单元格点击信号
        self.table.cellClicked.connect(self._on_cell_clicked)
    
    def _on_cell_clicked(self, row, column):
        """
        单元格点击事件处理
        
        Args:
            row (int): 行索引
            column (int): 列索引
        """
        # 发射信号 - 行是真实类别，列是预测类别
        self.cell_clicked.emit(row, column)
    
    def update_matrix(self, matrix, labels=None):
        """
        更新混淆矩阵
        
        Args:
            matrix (np.ndarray): 混淆矩阵
            labels (list, optional): 类别标签
        """
        if labels:
            self.labels = labels
        
        # 存储矩阵
        self.matrix = matrix
        
        # 设置表格大小
        n = len(self.labels) if self.labels else matrix.shape[0]
        self.table.setRowCount(n)
        self.table.setColumnCount(n)
        
        # 设置表头
        self.table.setHorizontalHeaderLabels(self.labels)
        self.table.setVerticalHeaderLabels(self.labels)
        
        # 填充数据
        for i in range(n):
            for j in range(n):
                value = matrix[i, j]
                item = QTableWidgetItem(str(value))
                
                # 为对角线元素设置背景色
                if i == j:
                    item.setBackground(QColor(200, 250, 200))  # 浅绿色
                
                # 将最大值高亮显示
                if value == np.max(matrix[i, :]):
                    item.setForeground(QColor(Qt.blue))
                
                self.table.setItem(i, j, item)


class TrainingMonitor(QWidget):
    """训练监控组件"""
    
    # 自定义信号
    start_training = pyqtSignal(dict)  # 开始训练信号，传递训练参数
    stop_training = pyqtSignal()  # 停止训练信号
    convert_to_annotation = pyqtSignal(str, str, list)  # 转换为标注任务信号 (真实类, 预测类, 样本列表)
    
    def __init__(self, parent=None):
        """初始化"""
        super().__init__(parent)
        
        # 内部状态
        self.is_training = False
        self.current_epoch = 0
        self.total_epochs = 0
        
        # 设置布局
        self._setup_ui()
    
    def _setup_ui(self):
        """设置界面"""
        # 主布局
        main_layout = QVBoxLayout(self)
        
        # 创建水平分割器
        splitter = QSplitter(Qt.Horizontal)
        main_layout.addWidget(splitter)
        
        # 左侧 - 参数设置和控制
        left_widget = QWidget()
        left_layout = QVBoxLayout(left_widget)
        
        # 参数设置组
        param_group = QGroupBox("训练参数")
        param_layout = QFormLayout()
        
        # 添加算法选择
        self.algorithm_combo = QComboBox()
        # 添加分类算法
        self.algorithm_combo.addItem("EfficientNet (平衡精度与速度)", "efficientnet")
        self.algorithm_combo.addItem("MobileNetV3 (轻量级)", "mobilenetv3")
        self.algorithm_combo.addItem("ResNet (高精度)", "resnet")
        self.algorithm_combo.addItem("Vision Transformer (ViT)", "vit")
        # 添加检测算法
        self.algorithm_combo.addItem("YOLOv8 (实时性好)", "yolov8")
        self.algorithm_combo.addItem("Faster R-CNN (高精度)", "fasterrcnn")
        self.algorithm_combo.addItem("NanoDet (超轻量级)", "nanodet")
        self.algorithm_combo.currentIndexChanged.connect(self.update_algorithm_advanced_options)
        param_layout.addRow("训练算法:", self.algorithm_combo)
        
        # 批量大小
        self.batch_size_spin = QSpinBox()
        self.batch_size_spin.setRange(1, 128)
        self.batch_size_spin.setValue(32)
        param_layout.addRow("批量大小:", self.batch_size_spin)
        
        # 学习率
        self.learning_rate_spin = QDoubleSpinBox()
        self.learning_rate_spin.setRange(0.0001, 0.1)
        self.learning_rate_spin.setSingleStep(0.0001)
        self.learning_rate_spin.setDecimals(4)
        self.learning_rate_spin.setValue(0.001)
        param_layout.addRow("学习率:", self.learning_rate_spin)
        
        # 训练轮数
        self.epochs_spin = QSpinBox()
        self.epochs_spin.setRange(1, 1000)
        self.epochs_spin.setValue(50)
        param_layout.addRow("训练轮数:", self.epochs_spin)
        
        # 测试集比例
        self.val_split_spin = QDoubleSpinBox()
        self.val_split_spin.setRange(0.1, 0.5)
        self.val_split_spin.setSingleStep(0.05)
        self.val_split_spin.setDecimals(2)
        self.val_split_spin.setValue(0.2)
        param_layout.addRow("测试集比例:", self.val_split_spin)
        
        # 早停
        self.early_stopping_check = QCheckBox()
        self.early_stopping_check.setChecked(True)
        param_layout.addRow("启用早停:", self.early_stopping_check)
        
        # 早停容忍度
        self.patience_spin = QSpinBox()
        self.patience_spin.setRange(1, 20)
        self.patience_spin.setValue(5)
        param_layout.addRow("早停容忍度:", self.patience_spin)
        
        # 设置布局
        param_group.setLayout(param_layout)
        
        # 添加高级参数组
        self.advanced_group = QGroupBox("高级参数")
        self.advanced_group.setCheckable(True)
        self.advanced_group.setChecked(False)
        self.advanced_layout = QFormLayout(self.advanced_group)
        
        # 创建常见高级参数控件
        self.weight_decay_spin = QDoubleSpinBox()
        self.weight_decay_spin.setRange(0.00001, 0.1)
        self.weight_decay_spin.setValue(0.0001)
        self.weight_decay_spin.setDecimals(5)
        self.weight_decay_spin.setSingleStep(0.00001)
        self.advanced_layout.addRow("权重衰减:", self.weight_decay_spin)
        
        self.momentum_spin = QDoubleSpinBox()
        self.momentum_spin.setRange(0.0, 1.0)
        self.momentum_spin.setValue(0.9)
        self.momentum_spin.setDecimals(2)
        self.momentum_spin.setSingleStep(0.05)
        self.advanced_layout.addRow("动量:", self.momentum_spin)
        
        # 预训练模型选项
        self.pretrained_check = QCheckBox()
        self.pretrained_check.setChecked(True)
        self.advanced_layout.addRow("使用预训练权重:", self.pretrained_check)
        
        # 数据增强选项
        self.augmentation_check = QCheckBox()
        self.augmentation_check.setChecked(True)
        self.advanced_layout.addRow("启用数据增强:", self.augmentation_check)
        
        # 添加到左侧布局
        left_layout.addWidget(param_group)
        left_layout.addWidget(self.advanced_group)
        
        # 添加特定算法的高级参数容器（初始为空）
        self.algorithm_specific_group = QGroupBox("算法特定参数")
        self.algorithm_specific_layout = QFormLayout(self.algorithm_specific_group)
        left_layout.addWidget(self.algorithm_specific_group)
        self.algorithm_specific_group.setVisible(False)
        
        # 初始化算法特定选项
        self.update_algorithm_advanced_options()
        
        # 进度显示
        progress_group = QGroupBox("训练进度")
        progress_layout = QVBoxLayout()
        
        # 当前轮次/总轮次
        self.epoch_label = QLabel("轮次: 0/0")
        progress_layout.addWidget(self.epoch_label)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        progress_layout.addWidget(self.progress_bar)
        
        # 当前损失和准确率
        metrics_layout = QFormLayout()
        self.train_loss_label = QLabel("N/A")
        self.val_loss_label = QLabel("N/A")
        self.train_acc_label = QLabel("N/A")
        self.val_acc_label = QLabel("N/A")
        
        metrics_layout.addRow("训练损失:", self.train_loss_label)
        metrics_layout.addRow("验证损失:", self.val_loss_label)
        metrics_layout.addRow("训练准确率:", self.train_acc_label)
        metrics_layout.addRow("验证准确率:", self.val_acc_label)
        
        progress_layout.addLayout(metrics_layout)
        progress_group.setLayout(progress_layout)
        left_layout.addWidget(progress_group)
        
        # 控制按钮组
        control_layout = QHBoxLayout()
        self.start_button = QPushButton("开始训练")
        self.start_button.clicked.connect(self._on_start_clicked)
        self.stop_button = QPushButton("停止训练")
        self.stop_button.clicked.connect(self._on_stop_clicked)
        self.stop_button.setEnabled(False)
        control_layout.addWidget(self.start_button)
        control_layout.addWidget(self.stop_button)
        
        left_layout.addLayout(control_layout)
        
        # 添加左侧部件到分割器
        splitter.addWidget(left_widget)
        
        # 右侧 - 图表和混淆矩阵
        right_widget = QWidget()
        right_layout = QVBoxLayout(right_widget)
        
        # 创建选项卡
        self.tab_widget = QTabWidget()
        
        # 图表选项卡
        self.chart = TrainingChart()
        self.tab_widget.addTab(self.chart, "训练图表")
        
        # 混淆矩阵选项卡
        self.confusion_matrix = ConfusionMatrix()
        self.confusion_matrix.cell_clicked.connect(self._on_confusion_cell_clicked)
        self.tab_widget.addTab(self.confusion_matrix, "混淆矩阵")
        
        # 添加样本列表部分到右侧混淆矩阵选项卡下方
        self.samples_group = QGroupBox("问题样本")
        samples_layout = QVBoxLayout()
        
        # 样本信息标签
        self.samples_info_label = QLabel("选择混淆矩阵单元格查看对应的问题样本")
        samples_layout.addWidget(self.samples_info_label)
        
        # 样本列表表格
        self.samples_table = QTableWidget(0, 3)  # 3列：样本ID, 真实标签, 预测标签
        self.samples_table.setHorizontalHeaderLabels(["样本ID", "真实标签", "预测标签"])
        self.samples_table.setEditTriggers(QTableWidget.NoEditTriggers)
        self.samples_table.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        samples_layout.addWidget(self.samples_table)
        
        # 创建批量转换为标注任务的按钮
        self.convert_button = QPushButton("转换为标注任务")
        self.convert_button.setEnabled(False)
        self.convert_button.clicked.connect(self._on_convert_clicked)
        samples_layout.addWidget(self.convert_button)
        
        self.samples_group.setLayout(samples_layout)
        
        # 添加到右侧布局
        right_layout.addWidget(self.tab_widget)
        right_layout.addWidget(self.samples_group)
        self.samples_group.setVisible(False)  # 初始不可见
        
        # 设置分割器比例
        splitter.setSizes([300, 700])
    
    def _on_start_clicked(self):
        """开始训练按钮点击事件"""
        # 获取训练参数
        params = self.get_training_params()
        
        # 更新UI状态
        self.is_training = True
        self.start_button.setEnabled(False)
        self.stop_button.setEnabled(True)
        self.batch_size_spin.setEnabled(False)
        self.learning_rate_spin.setEnabled(False)
        self.epochs_spin.setEnabled(False)
        self.val_split_spin.setEnabled(False)
        self.early_stopping_check.setEnabled(False)
        self.patience_spin.setEnabled(False)
        self.algorithm_combo.setEnabled(False)
        
        # 重置进度
        self.current_epoch = 0
        self.total_epochs = params['epochs']
        self.progress_bar.setValue(0)
        self.epoch_label.setText(f"轮次: 0/{self.total_epochs}")
        
        self.train_loss_label.setText("N/A")
        self.val_loss_label.setText("N/A")
        self.train_acc_label.setText("N/A")
        self.val_acc_label.setText("N/A")
        
        # 清空图表数据
        if hasattr(self.chart, 'train_loss'):
            self.chart.train_loss = []
            self.chart.val_loss = []
            self.chart.train_acc = []
            self.chart.val_acc = []
            self.chart.learning_rates = []
            self.chart.epochs = []
            self.chart.update_chart()
        
        # 发射开始训练信号
        self.start_training.emit(params)
    
    def _on_stop_clicked(self):
        """停止训练按钮点击事件"""
        # 更新UI状态
        self.is_training = False
        self.start_button.setEnabled(True)
        self.stop_button.setEnabled(False)
        self.batch_size_spin.setEnabled(True)
        self.learning_rate_spin.setEnabled(True)
        self.epochs_spin.setEnabled(True)
        self.val_split_spin.setEnabled(True)
        self.early_stopping_check.setEnabled(True)
        self.patience_spin.setEnabled(True)
        self.algorithm_combo.setEnabled(True)
        
        # 发射停止训练信号
        self.stop_training.emit()
    
    def update_progress(self, epoch, metrics):
        """
        更新训练进度
        
        Args:
            epoch (int): 当前轮次（从0开始）
            metrics (dict): 训练指标字典
        """
        # 保存指标数据的本地副本
        metrics_copy = metrics.copy() if metrics else {}
        current_epoch = epoch
        
        # 使用QTimer确保UI更新在主线程中进行
        def _update_ui():
            # 更新当前轮次
            self.current_epoch = current_epoch + 1  # epoch从0开始，显示从1开始
            
            # 更新进度条
            progress = int((self.current_epoch / self.total_epochs) * 100)
            self.progress_bar.setValue(progress)
            
            # 更新轮次标签
            self.epoch_label.setText(f"轮次: {self.current_epoch}/{self.total_epochs}")
            
            # 更新指标标签
            if 'train_loss' in metrics_copy:
                self.train_loss_label.setText(f"{metrics_copy['train_loss']:.4f}")
            
            if 'val_loss' in metrics_copy:
                self.val_loss_label.setText(f"{metrics_copy['val_loss']:.4f}")
            
            if 'train_acc' in metrics_copy:
                self.train_acc_label.setText(f"{metrics_copy['train_acc']:.2f}%")
            
            if 'val_acc' in metrics_copy:
                self.val_acc_label.setText(f"{metrics_copy['val_acc']:.2f}%")
            
            # 更新图表
            self.chart.update_data(current_epoch, metrics_copy)
        
        # 使用QTimer单次触发在主线程中更新UI
        QTimer.singleShot(0, _update_ui)
    
    def training_completed(self, final_metrics=None):
        """
        训练完成处理
        
        Args:
            final_metrics (dict, optional): 最终训练指标
        """
        # 保存指标数据的本地副本
        metrics_copy = final_metrics.copy() if final_metrics else {}
        
        # 使用QTimer确保UI更新在主线程中进行
        def _update_ui():
            # 更新UI状态
            self.is_training = False
            self.start_button.setEnabled(True)
            self.stop_button.setEnabled(False)
            self.batch_size_spin.setEnabled(True)
            self.learning_rate_spin.setEnabled(True)
            self.epochs_spin.setEnabled(True)
            self.val_split_spin.setEnabled(True)
            self.early_stopping_check.setEnabled(True)
            self.patience_spin.setEnabled(True)
            self.algorithm_combo.setEnabled(True)
            
            # 更新进度条到100%
            self.progress_bar.setValue(100)
            
            # 如果提供了混淆矩阵，则更新混淆矩阵选项卡
            if metrics_copy and 'confusion_matrix' in metrics_copy:
                labels = metrics_copy.get('labels', None)
                self.confusion_matrix.update_matrix(metrics_copy['confusion_matrix'], labels)
                
                # 切换到混淆矩阵选项卡
                self.tab_widget.setCurrentWidget(self.confusion_matrix)
        
        # 使用QTimer单次触发在主线程中更新UI
        QTimer.singleShot(0, _update_ui)
    
    def get_training_params(self):
        """
        获取当前训练参数
        
        Returns:
            dict: 训练参数字典
        """
        # 获取当前选择的算法
        algorithm_index = self.algorithm_combo.currentIndex()
        algorithm_data = self.algorithm_combo.itemData(algorithm_index)
        
        # 基本参数
        params = {
            'batch_size': self.batch_size_spin.value(),
            'learning_rate': self.learning_rate_spin.value(),
            'epochs': self.epochs_spin.value(),
            'val_split': self.val_split_spin.value(),
            'early_stopping': self.early_stopping_check.isChecked(),
            'patience': self.patience_spin.value(),
            'algorithm': algorithm_data
        }
        
        # 如果高级参数组被启用，添加高级参数
        if self.advanced_group.isChecked():
            params.update({
                'weight_decay': self.weight_decay_spin.value(),
                'momentum': self.momentum_spin.value(),
                'pretrained': self.pretrained_check.isChecked(),
                'data_augmentation': self.augmentation_check.isChecked()
            })
            
            # 添加算法特定参数
            if self.algorithm_specific_group.isVisible():
                # 获取算法特定组件
                special_params = {}
                
                for i in range(self.algorithm_specific_layout.rowCount()):
                    # 获取标签和对应的控件
                    label_item = self.algorithm_specific_layout.itemAt(i*2)
                    field_item = self.algorithm_specific_layout.itemAt(i*2+1)
                    
                    if label_item and field_item and label_item.widget() and field_item.widget():
                        label = label_item.widget().text().replace(":", "")
                        widget = field_item.widget()
                        
                        # 根据控件类型获取值
                        if isinstance(widget, QComboBox):
                            special_params[label] = widget.currentText()
                        elif isinstance(widget, QSpinBox) or isinstance(widget, QDoubleSpinBox):
                            special_params[label] = widget.value()
                        elif isinstance(widget, QCheckBox):
                            special_params[label] = widget.isChecked()
                
                # 更新参数
                if special_params:
                    params['special_params'] = special_params
                
                # 对特定算法做特殊处理（如需要转换参数格式等）
                if algorithm_data == "efficientnet":
                    if "模型变体" in special_params:
                        params['model_variant'] = special_params["模型变体"]
                    if "宽度倍数" in special_params:
                        params['width_multiplier'] = special_params["宽度倍数"]
                
                elif algorithm_data == "mobilenetv3":
                    if "模型大小" in special_params:
                        params['model_size'] = special_params["模型大小"].lower()
                    if "宽度倍数" in special_params:
                        params['alpha'] = float(special_params["宽度倍数"])
                
                elif algorithm_data == "resnet":
                    if "层数" in special_params:
                        # 从"ResNet50"中提取数字部分
                        layer_str = special_params["层数"]
                        if "ResNet" in layer_str:
                            params['layers'] = int(layer_str.replace("ResNet", ""))
                
                elif algorithm_data == "vit":
                    if "补丁大小" in special_params:
                        params['patch_size'] = int(special_params["补丁大小"])
                    if "注意力头数" in special_params:
                        params['num_heads'] = special_params["注意力头数"]
                    if "模型大小" in special_params:
                        params['model_size'] = special_params["模型大小"].lower()
                
                elif algorithm_data == "yolov8":
                    if "模型大小" in special_params:
                        params['model_size'] = special_params["模型大小"]
                    if "置信度阈值" in special_params:
                        params['conf_thres'] = special_params["置信度阈值"]
                
                elif algorithm_data == "fasterrcnn" or algorithm_data == "retinanet" or algorithm_data == "ssd" or algorithm_data == "deeplabv3":
                    if "主干网络" in special_params:
                        params['backbone'] = special_params["主干网络"]
                        
                elif algorithm_data == "sam":
                    if "模型类型" in special_params:
                        params['model_type'] = special_params["模型类型"]
                    if "格点数量" in special_params:
                        params['points_per_side'] = special_params["格点数量"]
                        
                elif algorithm_data == "unet":
                    if "起始滤波器" in special_params:
                        params['filters'] = special_params["起始滤波器"]
                    if "网络深度" in special_params:
                        params['depth'] = special_params["网络深度"]
        
        return params
    
    def _on_confusion_cell_clicked(self, true_class, pred_class):
        """
        混淆矩阵单元格点击事件处理
        
        Args:
            true_class (int): 真实类别索引
            pred_class (int): 预测类别索引
        """
        # 显示样本组
        self.samples_group.setVisible(True)
        
        # 获取类别名称
        true_label = self.confusion_matrix.labels[true_class]
        pred_label = self.confusion_matrix.labels[pred_class]
        
        # 更新样本信息标签
        if true_class == pred_class:
            self.samples_info_label.setText(f"类别 [{true_label}] 的正确分类样本")
            # 禁用转换按钮
            self.convert_button.setEnabled(False)
        else:
            self.samples_info_label.setText(f"真实类别 [{true_label}] 被错误分类为 [{pred_label}] 的样本")
            # 启用转换按钮
            self.convert_button.setEnabled(True)
        
        # 存储当前选中的类别
        self.current_true_class = true_label
        self.current_pred_class = pred_label
        
        # 更新样本表格（假设我们有问题样本的ID列表）
        # 这里需要模型训练器提供具体样本数据
        if hasattr(self, 'misclassified_samples') and (true_class, pred_class) in self.misclassified_samples:
            samples = self.misclassified_samples[(true_class, pred_class)]
            self._update_samples_table(samples, true_label, pred_label)
        else:
            # 清空表格
            self.samples_table.setRowCount(0)
            self.convert_button.setEnabled(False)
    
    def _update_samples_table(self, samples, true_label, pred_label):
        """
        更新样本表格
        
        Args:
            samples (list): 样本ID列表
            true_label (str): 真实类别名称
            pred_label (str): 预测类别名称
        """
        # 清空表格
        self.samples_table.setRowCount(0)
        
        # 添加样本
        self.samples_table.setRowCount(len(samples))
        for i, sample_id in enumerate(samples):
            # 样本ID
            id_item = QTableWidgetItem(str(sample_id))
            self.samples_table.setItem(i, 0, id_item)
            
            # 真实标签
            true_item = QTableWidgetItem(true_label)
            self.samples_table.setItem(i, 1, true_item)
            
            # 预测标签
            pred_item = QTableWidgetItem(pred_label)
            self.samples_table.setItem(i, 2, pred_item)
    
    def _on_convert_clicked(self):
        """转换为标注任务按钮点击事件"""
        # 获取当前选中的样本
        selected_rows = set()
        for item in self.samples_table.selectedItems():
            selected_rows.add(item.row())
        
        # 如果没有选中任何行，则使用所有行
        if not selected_rows:
            selected_rows = range(self.samples_table.rowCount())
        
        # 获取选中的样本ID
        samples = []
        for row in selected_rows:
            sample_id = self.samples_table.item(row, 0).text()
            samples.append(sample_id)
        
        # 发射转换信号
        if samples:
            self.convert_to_annotation.emit(self.current_true_class, self.current_pred_class, samples)
    
    def set_misclassified_samples(self, samples_dict):
        """
        设置错分类样本数据
        
        Args:
            samples_dict (dict): 键为(true_class_idx, pred_class_idx)，值为样本ID列表
        """
        self.misclassified_samples = samples_dict

    def filter_algorithms_by_type(self, dataset_type):
        """
        根据数据集类型过滤算法选择
        
        Args:
            dataset_type (str): 数据集类型，如 'classification', 'detection'
        """
        # 清空当前选择
        current_index = self.algorithm_combo.currentIndex()
        current_data = self.algorithm_combo.itemData(current_index) if current_index >= 0 else None
        
        # 阻止事件触发
        self.algorithm_combo.blockSignals(True)
        self.algorithm_combo.clear()
        
        if dataset_type == "classification":
            # 添加分类算法 - 从TorchVision/Keras
            self.algorithm_combo.addItem("EfficientNet (平衡精度与速度)", "efficientnet")
            self.algorithm_combo.addItem("MobileNetV3 (轻量级)", "mobilenetv3")
            self.algorithm_combo.addItem("ResNet (高精度)", "resnet")
            self.algorithm_combo.addItem("Vision Transformer (ViT)", "vit")
            self.algorithm_combo.addItem("Xception (高级特征提取)", "xception")
            self.algorithm_combo.addItem("DenseNet (密集连接)", "densenet")
            self.algorithm_combo.addItem("InceptionV3 (Google)", "inceptionv3")
            self.algorithm_combo.addItem("ConvNeXt (最新卷积)", "convnext")
        elif dataset_type == "detection":
            # 添加检测算法
            self.algorithm_combo.addItem("YOLOv8 (实时性好)", "yolov8")
            self.algorithm_combo.addItem("Faster R-CNN (高精度)", "fasterrcnn")
            self.algorithm_combo.addItem("RetinaNet (单阶段高精度)", "retinanet")
            self.algorithm_combo.addItem("NanoDet (超轻量级)", "nanodet")
            self.algorithm_combo.addItem("CenterNet (关键点)", "centernet")
            self.algorithm_combo.addItem("DETR (基于Transformer)", "detr")
            self.algorithm_combo.addItem("SSD (单发多框检测)", "ssd")
            self.algorithm_combo.addItem("EfficientDet (Google)", "efficientdet")
        elif dataset_type == "segmentation":
            # 添加分割算法
            self.algorithm_combo.addItem("DeepLabV3+ (语义分割)", "deeplabv3")
            self.algorithm_combo.addItem("U-Net (医学图像分割)", "unet")
            self.algorithm_combo.addItem("Mask R-CNN (实例分割)", "maskrcnn")
            self.algorithm_combo.addItem("FCN (全卷积网络)", "fcn")
            self.algorithm_combo.addItem("SegFormer (Transformer分割)", "segformer")
            self.algorithm_combo.addItem("SAM (分割万物模型)", "sam")
            self.algorithm_combo.addItem("PanopticFPN (全景分割)", "panopticfpn")
        else:
            # 未知类型，添加通用算法
            self.algorithm_combo.addItem("通用模型", "generic")
        
        # 恢复事件触发
        self.algorithm_combo.blockSignals(False)
        
        # 尝试恢复之前选择的算法类型
        if current_data:
            index = self.algorithm_combo.findData(current_data)
            if index >= 0:
                self.algorithm_combo.setCurrentIndex(index)
            elif self.algorithm_combo.count() > 0:
                self.algorithm_combo.setCurrentIndex(0)
                
    def update_algorithm_advanced_options(self):
        """根据当前选择的算法更新高级选项"""
        # 获取当前选择的算法
        current_index = self.algorithm_combo.currentIndex()
        algorithm = self.algorithm_combo.itemData(current_index) if current_index >= 0 else None
        
        # 清除现有的算法特定选项
        while self.algorithm_specific_layout.count() > 0:
            item = self.algorithm_specific_layout.takeAt(0)
            if item.widget():
                item.widget().deleteLater()
        
        if not algorithm:
            self.algorithm_specific_group.setVisible(False)
            return
            
        # 根据算法类型显示特定选项
        self.algorithm_specific_group.setVisible(True)
        
        if algorithm == "efficientnet":
            # EfficientNet特定选项
            model_variant = QComboBox()
            model_variant.addItems(["B0", "B1", "B2", "B3", "B4", "B5", "B6", "B7"])
            model_variant.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("模型变体:", model_variant)
            
            # 添加额外选项
            width_mult = QDoubleSpinBox()
            width_mult.setRange(0.5, 2.0)
            width_mult.setValue(1.0)
            width_mult.setSingleStep(0.1)
            width_mult.setDecimals(1)
            self.algorithm_specific_layout.addRow("宽度倍数:", width_mult)
            
        elif algorithm == "mobilenetv3":
            # MobileNetV3特定选项
            model_size = QComboBox()
            model_size.addItems(["Small", "Large"])
            model_size.setCurrentIndex(1)  # Large
            self.algorithm_specific_layout.addRow("模型大小:", model_size)
            
            alpha = QComboBox()
            alpha.addItems(["0.75", "1.0", "1.25"])
            alpha.setCurrentIndex(1)  # 1.0
            self.algorithm_specific_layout.addRow("宽度倍数:", alpha)
            
        elif algorithm == "resnet":
            # ResNet特定选项
            model_layers = QComboBox()
            model_layers.addItems(["ResNet18", "ResNet34", "ResNet50", "ResNet101", "ResNet152"])
            model_layers.setCurrentIndex(2)  # ResNet50
            self.algorithm_specific_layout.addRow("层数:", model_layers)
            
        elif algorithm == "vit":
            # Vision Transformer特定选项
            patch_size = QComboBox()
            patch_size.addItems(["8", "16", "32"])
            patch_size.setCurrentIndex(1)  # 16
            self.algorithm_specific_layout.addRow("补丁大小:", patch_size)
            
            heads = QSpinBox()
            heads.setRange(1, 16)
            heads.setValue(8)
            self.algorithm_specific_layout.addRow("注意力头数:", heads)
            
            model_size = QComboBox()
            model_size.addItems(["Base", "Large", "Huge"])
            model_size.setCurrentIndex(0)  # Base
            self.algorithm_specific_layout.addRow("模型大小:", model_size)
            
        elif algorithm == "yolov8":
            # YOLOv8特定选项
            model_size = QComboBox()
            model_size.addItems(["nano", "small", "medium", "large", "xlarge"])
            model_size.setCurrentIndex(2)  # medium
            self.algorithm_specific_layout.addRow("模型大小:", model_size)
            
            conf_thres = QDoubleSpinBox()
            conf_thres.setRange(0.1, 0.9)
            conf_thres.setValue(0.25)
            conf_thres.setSingleStep(0.05)
            self.algorithm_specific_layout.addRow("置信度阈值:", conf_thres)
            
        elif algorithm == "fasterrcnn":
            # Faster R-CNN特定选项
            backbone = QComboBox()
            backbone.addItems(["ResNet50", "ResNet101", "MobileNetV3"])
            backbone.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("主干网络:", backbone)
            
            rpn_pre_nms = QSpinBox()
            rpn_pre_nms.setRange(100, 2000)
            rpn_pre_nms.setValue(1000)
            rpn_pre_nms.setSingleStep(100)
            self.algorithm_specific_layout.addRow("RPN提议数:", rpn_pre_nms)
            
        elif algorithm == "retinanet":
            # RetinaNet特定选项
            backbone = QComboBox()
            backbone.addItems(["ResNet50", "ResNet101", "ResNeXt101"])
            backbone.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("主干网络:", backbone)
            
            focal_alpha = QDoubleSpinBox()
            focal_alpha.setRange(0.1, 0.9)
            focal_alpha.setValue(0.25)
            focal_alpha.setSingleStep(0.05)
            self.algorithm_specific_layout.addRow("Focal Alpha:", focal_alpha)
            
            focal_gamma = QDoubleSpinBox()
            focal_gamma.setRange(0.5, 5.0)
            focal_gamma.setValue(2.0)
            focal_gamma.setSingleStep(0.5)
            self.algorithm_specific_layout.addRow("Focal Gamma:", focal_gamma)
            
        elif algorithm == "ssd":
            # SSD特定选项
            backbone = QComboBox()
            backbone.addItems(["VGG16", "MobileNetV2", "ResNet50"])
            backbone.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("主干网络:", backbone)
            
            aspect_ratios = QComboBox()
            aspect_ratios.addItems(["[1.0, 2.0, 0.5]", "[1.0, 2.0, 3.0, 0.5, 0.33]"])
            aspect_ratios.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("长宽比:", aspect_ratios)
        
        elif algorithm == "deeplabv3":
            # DeepLabV3特定选项
            backbone = QComboBox()
            backbone.addItems(["ResNet50", "ResNet101", "MobileNetV3"])
            backbone.setCurrentIndex(1)  # ResNet101
            self.algorithm_specific_layout.addRow("主干网络:", backbone)
            
            aspp_rates = QComboBox()
            aspp_rates.addItems(["[6, 12, 18]", "[12, 24, 36]"])
            aspp_rates.setCurrentIndex(0)
            self.algorithm_specific_layout.addRow("ASPP率:", aspp_rates)
            
        elif algorithm == "unet":
            # U-Net特定选项
            filters = QSpinBox()
            filters.setRange(16, 128)
            filters.setValue(64)
            filters.setSingleStep(16)
            self.algorithm_specific_layout.addRow("起始滤波器:", filters)
            
            depth = QSpinBox()
            depth.setRange(3, 6)
            depth.setValue(4)
            self.algorithm_specific_layout.addRow("网络深度:", depth)
            
        elif algorithm == "sam":
            # SAM特定选项
            model_type = QComboBox()
            model_type.addItems(["vit_b", "vit_l", "vit_h"])
            model_type.setCurrentIndex(0)  # vit_b
            self.algorithm_specific_layout.addRow("模型类型:", model_type)
            
            points_per_side = QSpinBox()
            points_per_side.setRange(16, 64)
            points_per_side.setValue(32)
            points_per_side.setSingleStep(8)
            self.algorithm_specific_layout.addRow("格点数量:", points_per_side)
            
        elif algorithm == "xception":
            # Xception特定选项
            depth_multiplier = QDoubleSpinBox()
            depth_multiplier.setRange(0.5, 2.0)
            depth_multiplier.setValue(1.0)
            depth_multiplier.setSingleStep(0.1)
            depth_multiplier.setDecimals(1)
            self.algorithm_specific_layout.addRow("深度倍数:", depth_multiplier)
            
        elif algorithm == "densenet":
            # DenseNet特定选项
            model_type = QComboBox()
            model_type.addItems(["DenseNet121", "DenseNet169", "DenseNet201"])
            model_type.setCurrentIndex(0)  # DenseNet121
            self.algorithm_specific_layout.addRow("模型变体:", model_type)
            
            growth_rate = QSpinBox()
            growth_rate.setRange(12, 48)
            growth_rate.setValue(32)
            growth_rate.setSingleStep(4)
            self.algorithm_specific_layout.addRow("增长率:", growth_rate)
            
        # 其他算法可以类似添加 